/-
Copyright (c) 2025 Kexing Ying. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Kexing Ying
-/
module

public import Mathlib.Probability.Process.Adapted

/-!
# Predictable σ-algebra

This file defines the predictable σ-algebra associated to a filtration, as well as the
notion of predictable processes. We prove that predictable processes are progressively measurable
and adapted. We also give an equivalent characterization of predictability for discrete processes.

## Main definitions

* `Filtration.predictable` : The predictable σ-algebra associated to a filtration.
* `IsPredictable` : A process is predictable if it is measurable with respect to the
  predictable σ-algebra.

## Main results

* `IsPredictable.progMeasurable` : A predictable process is progressively measurable.
* `isPredictable_iff_measurable_add_one` : `u` is a discrete predictable process iff
  `u (n + 1)` is `𝓕 n`-measurable and `u 0` is `𝓕 0`-measurable.

## Tags

predictable, previsible

-/

@[expose] public section

open Filter Order TopologicalSpace

open scoped MeasureTheory NNReal ENNReal Topology

namespace MeasureTheory

variable {Ω ι : Type*} {m : MeasurableSpace Ω} {E : Type*} [TopologicalSpace E]

section

variable [Preorder ι] [OrderBot ι]

namespace Filtration

/-- Given a filtration `𝓕`, the predictable σ-algebra is the σ-algebra on `ι × Ω` generated by
sets of the form `(t, ∞) × A` for `t ∈ ι` and `A ∈ 𝓕 t` and `{⊥} × A` for `A ∈ 𝓕 ⊥`. -/
def predictable (𝓕 : Filtration ι m) : MeasurableSpace (ι × Ω) :=
  MeasurableSpace.generateFrom <|
    {s | ∃ A, MeasurableSet[𝓕 ⊥] A ∧ s = {⊥} ×ˢ A} ∪
    {s | ∃ i A, MeasurableSet[𝓕 i] A ∧ s = Set.Ioi i ×ˢ A}

end Filtration

/-- A process is said to be predictable if it is measurable with respect to the predictable
σ-algebra. -/
def IsPredictable (𝓕 : Filtration ι m) (u : ι → Ω → E) :=
  StronglyMeasurable[𝓕.predictable] <| Function.uncurry u

end

lemma measurableSet_predictable_singleton_bot_prod [LinearOrder ι] [OrderBot ι]
    {𝓕 : Filtration ι m} {s : Set Ω} (hs : MeasurableSet[𝓕 ⊥] s) :
    MeasurableSet[𝓕.predictable] <| {⊥} ×ˢ s :=
  MeasurableSpace.measurableSet_generateFrom <| Or.inl ⟨s, hs, rfl⟩

lemma measurableSet_predictable_Ioi_prod [LinearOrder ι] [OrderBot ι]
    {𝓕 : Filtration ι m} {i : ι} {s : Set Ω} (hs : MeasurableSet[𝓕 i] s) :
    MeasurableSet[𝓕.predictable] <| Set.Ioi i ×ˢ s :=
  MeasurableSpace.measurableSet_generateFrom <| Or.inr ⟨i, s, hs, rfl⟩

/-- Sets of the form `(i, j] × A` for any `A ∈ 𝓕 i` are measurable with respect to the predictable
σ-algebra. -/
lemma measurableSet_predictable_Ioc_prod [LinearOrder ι] [OrderBot ι]
    {𝓕 : Filtration ι m} (i j : ι) {s : Set Ω} (hs : MeasurableSet[𝓕 i] s) :
    MeasurableSet[𝓕.predictable] <| Set.Ioc i j ×ˢ s := by
  obtain hij | hij := le_total j i
  · simp [hij]
  · rw [← Set.Ioi_diff_Ioi, (by simp : (Set.Ioi i \ Set.Ioi j) ×ˢ s
      = Set.Ioi i ×ˢ (s \ s) ∪ (Set.Ioi i \ Set.Ioi j) ×ˢ s), ← Set.prod_diff_prod]
    exact (measurableSet_predictable_Ioi_prod hs).diff
      (measurableSet_predictable_Ioi_prod <| 𝓕.mono hij _ hs)

lemma measurableSpace_le_predictable_of_measurableSet [Preorder ι] [OrderBot ι]
    {𝓕 : Filtration ι m} {m' : MeasurableSpace (ι × Ω)}
    (hm'bot : ∀ A, MeasurableSet[𝓕 ⊥] A → MeasurableSet[m'] ({⊥} ×ˢ A))
    (hm' : ∀ i A, MeasurableSet[𝓕 i] A → MeasurableSet[m'] ((Set.Ioi i) ×ˢ A)) :
    𝓕.predictable ≤ m' := by
  refine MeasurableSpace.generateFrom_le ?_
  rintro - (⟨A, hA, rfl⟩ | ⟨i, A, hA, rfl⟩)
  · exact hm'bot A hA
  · exact hm' i A hA

namespace IsPredictable

open Filtration

variable [LinearOrder ι] [OrderBot ι] [MeasurableSpace ι] [TopologicalSpace ι]
    [OpensMeasurableSpace ι] [OrderClosedTopology ι]
    [MetrizableSpace E] [MeasurableSpace E] [BorelSpace E] [SecondCountableTopology E]

/-- A predictable process is progressively measurable. -/
lemma progMeasurable {𝓕 : Filtration ι m} {u : ι → Ω → E} (h𝓕 : IsPredictable 𝓕 u) :
    ProgMeasurable 𝓕 u := by
  refine fun i ↦ Measurable.stronglyMeasurable ?_
  rw [IsPredictable, stronglyMeasurable_iff_measurable, measurable_iff_comap_le] at h𝓕
  rw [measurable_iff_comap_le, (by aesop : (fun (p : Set.Iic i × Ω) ↦ u (p.1) p.2)
      = Function.uncurry u ∘ (fun p ↦ (p.1, p.2))), ← MeasurableSpace.comap_comp]
  refine (MeasurableSpace.comap_mono h𝓕).trans <| MeasurableSpace.comap_le_iff_le_map.2 <|
    measurableSpace_le_predictable_of_measurableSet ?_ ?_
  · intros A hA
    simp only [MeasurableSpace.map_def,
      (by aesop : (fun (p : Set.Iic i × Ω) ↦ ((p.1 : ι), p.2)) ⁻¹' ({⊥} ×ˢ A) = {⊥} ×ˢ A)]
    exact (measurableSet_singleton _).prod <| 𝓕.mono bot_le _ hA
  · intros j A hA
    simp only [MeasurableSpace.map_def]
    obtain hji | hij := le_total j i
    · rw [(by grind : (fun (p : Set.Iic i × Ω) ↦ ((p.1 : ι), p.2)) ⁻¹' Set.Ioi j ×ˢ A
        = (Subtype.val ⁻¹' (Set.Ioc j i)) ×ˢ A)]
      exact (measurable_subtype_coe measurableSet_Ioc).prod (𝓕.mono hji _ hA)
    · simp [(by grind : (fun (p : Set.Iic i × Ω) ↦ ((p.1 : ι), p.2)) ⁻¹' Set.Ioi j ×ˢ A = ∅)]

/-- A predictable process is adapted. -/
lemma adapted {𝓕 : Filtration ι m} {u : ι → Ω → E} (h𝓕 : IsPredictable 𝓕 u) :
    Adapted 𝓕 u :=
  h𝓕.progMeasurable.adapted

omit [SecondCountableTopology E] in
lemma measurableSet_prodMk_add_one_of_predictable {𝓕 : Filtration ℕ m} {s : Set (ℕ × Ω)}
    (hs : MeasurableSet[𝓕.predictable] s) (n : ℕ) :
    MeasurableSet[𝓕 n] {ω | (n + 1, ω) ∈ s} := by
  rw [(by aesop : {ω | (n + 1, ω) ∈ s} = (Prod.mk (α := Set.singleton (n + 1)) (β := Ω)
      ⟨n + 1, rfl⟩) ⁻¹' ((fun (p : Set.singleton (n + 1) × Ω) ↦ ((p.1 : ℕ), p.2)) ⁻¹' s))]
  refine measurableSet_preimage (mβ := Subtype.instMeasurableSpace.prod (𝓕 n))
    measurable_prodMk_left <| measurableSet_preimage ?_ hs
  rw [measurable_iff_comap_le, MeasurableSpace.comap_le_iff_le_map]
  refine MeasurableSpace.generateFrom_le ?_
  rintro - (⟨A, hA, rfl⟩ | ⟨i, A, hA, rfl⟩)
  · rw [MeasurableSpace.map_def,
      (_ : (fun (p : Set.singleton (n + 1) × Ω) ↦ ((p.1 : ℕ), p.2)) ⁻¹' ({⊥} ×ˢ A) = ∅)]
    · simp
    · ext p
      simp only [Nat.bot_eq_zero, Set.mem_preimage, Set.mem_prod, Set.mem_singleton_iff,
        Set.mem_empty_iff_false, iff_false, not_and]
      exact fun hp1 ↦ False.elim <| Nat.succ_ne_zero n (hp1 ▸ p.1.2.symm)
  · rw [MeasurableSpace.map_def]
    obtain hni | hin := lt_or_ge n i
    · rw [(_ : (fun (p : Set.singleton (n + 1) × Ω) ↦ ((p.1 : ℕ), p.2)) ⁻¹' (Set.Ioi i ×ˢ A) = ∅)]
      · simp
      · ext p
        simp only [Set.mem_preimage, Set.mem_prod, Set.mem_Ioi, Set.mem_empty_iff_false,
          iff_false, not_and]
        rw [p.1.2]
        grind
    · rw [(_ : (fun (p : Set.singleton (n + 1) × Ω) ↦ ((p.1 : ℕ), p.2)) ⁻¹' (Set.Ioi i ×ˢ A)
          = {⟨n + 1, rfl⟩} ×ˢ A)]
      · exact MeasurableSet.prod (MeasurableSet.of_subtype_image trivial) (𝓕.mono hin _ hA)
      · ext p
        simp only [Set.mem_preimage, Set.mem_prod, Set.mem_Ioi, Set.mem_singleton_iff,
          and_congr_left_iff]
        intro hp2
        rw [p.1.2]
        exact ⟨fun _ ↦ by aesop, fun _ ↦ lt_add_one_iff.2 hin⟩

omit [SecondCountableTopology E] in
/-- If `u` is a discrete predictable process, then `u (n + 1)` is `𝓕 n`-measurable. -/
lemma measurable_add_one {𝓕 : Filtration ℕ m} {u : ℕ → Ω → E} (h𝓕 : IsPredictable 𝓕 u) (n : ℕ) :
    Measurable[𝓕 n] (u (n + 1)) := by
  intro s hs
  rw [(by aesop : u (n + 1) ⁻¹' s = {ω | (n + 1, ω) ∈ (Function.uncurry u) ⁻¹' s})]
  exact measurableSet_prodMk_add_one_of_predictable (h𝓕.measurable hs) n

end IsPredictable

section

variable [MetrizableSpace E] [MeasurableSpace E] [BorelSpace E]

lemma measurableSet_predictable_singleton_prod
    {𝓕 : Filtration ℕ m} {n : ℕ} {s : Set Ω} (hs : MeasurableSet[𝓕 n] s) :
    MeasurableSet[𝓕.predictable] <| {n + 1} ×ˢ s := by
  rw [(_ : {n + 1} = Set.Ioc n (n + 1))]
  · exact measurableSet_predictable_Ioc_prod _ _ hs
  · ext m
    simp only [Set.mem_singleton_iff, Set.mem_Ioc]
    omega

lemma isPredictable_of_measurable_add_one [SecondCountableTopology E]
    {𝓕 : Filtration ℕ m} {u : ℕ → Ω → E}
    (h₀ : Measurable[𝓕 0] (u 0)) (h : ∀ n, Measurable[𝓕 n] (u (n + 1))) :
    IsPredictable 𝓕 u := by
  refine Measurable.stronglyMeasurable ?_
  intro s hs
  rw [(by aesop : Function.uncurry u ⁻¹' s = ⋃ n : ℕ, {n} ×ˢ (u n ⁻¹' s))]
  refine MeasurableSet.iUnion <| fun n ↦ ?_
  obtain (rfl | hn) := n.eq_zero_or_eq_succ_pred
  · exact MeasurableSpace.measurableSet_generateFrom <| Or.inl ⟨u 0 ⁻¹' s, h₀ hs, rfl⟩
  · rw [hn]
    exact measurableSet_predictable_singleton_prod (h (n - 1) hs)

/-- A discrete process `u` is predictable iff `u (n + 1)` is `𝓕 n`-measurable for all `n` and
`u 0` is `𝓕 0`-measurable. -/
lemma isPredictable_iff_measurable_add_one [SecondCountableTopology E]
    {𝓕 : Filtration ℕ m} {u : ℕ → Ω → E} :
    IsPredictable 𝓕 u ↔ Measurable[𝓕 0] (u 0) ∧ ∀ n, Measurable[𝓕 n] (u (n + 1)) :=
  ⟨fun h𝓕 ↦ ⟨(h𝓕.adapted 0).measurable, fun n ↦ h𝓕.measurable_add_one (n)⟩,
    fun h ↦ isPredictable_of_measurable_add_one h.1 h.2⟩

end

end MeasureTheory
